package gjj.algorithm.ml;

import gjj.App;
import gjj.enums.ImgLabel;
import gjj.util.ImageUtils;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.MatOfKeyPoint;
import org.opencv.core.Size;
import org.opencv.features2d.BOWImgDescriptorExtractor;
import org.opencv.features2d.BOWKMeansTrainer;
import org.opencv.features2d.DescriptorMatcher;
import org.opencv.features2d.FlannBasedMatcher;
import org.opencv.features2d.SIFT;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.opencv.ml.Ml;
import org.opencv.ml.SVM;
import org.opencv.ml.TrainData;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class BowClassificationTrainer {

    private static final String ROOT = App.class.getResource("/").getPath();

    static {
        System.load(ROOT + "libopencv_java4120.dylib");
    }

    private static final int DICTIONARY_SIZE = 50;          // 视觉词个数（样本少就设小）
    private static final int DETECTOR_LIMIT = 300;          // 每张图最多提多少特征
    private static final double C = 2, GAMMA = 0.8;      // SVM 参数
    private static final Size UNIFORM_SIZE = new Size(224, 224);

    private static Mat bowExtract(Mat img, MatOfKeyPoint kp, BOWImgDescriptorExtractor bowExtractor) {
        Mat hist = new Mat();
        bowExtractor.compute(img, kp, hist);
        if (hist.empty()) hist = Mat.zeros(1, DICTIONARY_SIZE, CvType.CV_32F);
        hist.convertTo(hist, CvType.CV_32F);
        return hist;
    }

    public static void main(String[] args) {
        String rootPath = "/Users/kwok/Downloads/img_classification_train_data";
        List<Integer> trainLabels = new ArrayList<>();
        List<Mat> trainImages = new ArrayList<>();
        for (ImgLabel imgLabel : ImgLabel.values()) {
            File dir = new File(rootPath + "/" + imgLabel.getName());
            File[] dirFiles = dir.listFiles();
            if (!dir.exists() || null == dirFiles) {
                System.out.printf("文件夹【%s】不存在\n", dir.getAbsolutePath());
                continue;
            }
            System.out.printf("找到训练数据《%s》，开始遍历下面的文件...\n", imgLabel.getName());
            for (File file : dirFiles) {
                if (null == file) continue;
                System.out.printf("读取图片《%s》...\n", file.getName());
                Mat img = Imgcodecs.imread(file.getAbsolutePath());
                if (img.empty()) continue;
                img.convertTo(img, CvType.CV_8U);
                ImageUtils.adaptImgChannel(img);
                Imgproc.resize(img, img, UNIFORM_SIZE);
                Imgproc.GaussianBlur(img, img, new Size(5, 5), 0);
                trainImages.add(img);
                trainLabels.add(imgLabel.getLabel());
            }
        }
        // 1. 生成视觉词典
        SIFT detector = SIFT.create(DETECTOR_LIMIT);
        BOWKMeansTrainer bowTrainer = new BOWKMeansTrainer(DICTIONARY_SIZE);
        List<MatOfKeyPoint> trainImageKeypointList = new ArrayList<>(trainImages.size());
        for (Mat img : trainImages) {
            MatOfKeyPoint keypoint = new MatOfKeyPoint();
            Mat descriptors = new Mat();
            detector.detectAndCompute(img, new Mat(), keypoint, descriptors);
            if (keypoint.rows() == 0) continue;
            keypoint.convertTo(keypoint, CvType.CV_32F);
            descriptors.convertTo(descriptors, CvType.CV_32F);
            bowTrainer.add(descriptors);
            trainImageKeypointList.add(keypoint);
        }
        Mat dictionary = bowTrainer.cluster();
        dictionary.convertTo(dictionary, CvType.CV_32F);
        // 2. BOW 提取器
        BOWImgDescriptorExtractor bowExtractor = new BOWImgDescriptorExtractor(detector, FlannBasedMatcher.create(DescriptorMatcher.FLANNBASED));
        bowExtractor.setVocabulary(dictionary);
        // 3. 提取训练集直方图
        Mat trainMat = new Mat();
        Mat labelMat = new Mat(trainLabels.size(), 1, CvType.CV_32SC1);
        for (int i = 0; i < trainImages.size(); i++) {
            Mat hist = bowExtract(trainImages.get(i), trainImageKeypointList.get(i), bowExtractor);
            trainMat.push_back(hist);
            labelMat.put(i, 0, trainLabels.get(i));
        }
        // 4. 训练 SVM
        TrainData trainData = TrainData.create(trainMat, Ml.ROW_SAMPLE, labelMat);
        // 创建SVM对象并设置参数
        System.out.println("训练数据已经准备好，开始训练模型...");
        SVM svm = SVM.create();
        svm.setType(SVM.C_SVC); // 设置SVM类型为C_SVC（分类）
        svm.setKernel(SVM.RBF); // 设置核函数
        svm.setC(C);
        svm.setGamma(GAMMA);
//        svm.setTermCriteria(new TermCriteria(TermCriteria.MAX_ITER, 100, 1e-6));
        svm.train(trainData);
        //保存SVM模型
        System.out.println("训练完成，保存模型...");
        svm.save(ROOT + "svm_classification_bow.xml");
        // 5. 简单用训练测试测试一下
        int hit = 0;
        for (int i = 0; i < trainImages.size(); i++) {
            Mat hist = bowExtract(trainImages.get(i), trainImageKeypointList.get(i), bowExtractor);
            int predict = (int) svm.predict(hist);
            if (predict == trainLabels.get(i)) {
                hit++;
            }
        }
        System.out.printf("测试准确率: %.2f %% (%d/%d)\n", 100.0 * hit / trainImages.size(), hit, trainImages.size());
    }
}
